import axios from 'axios';

export class ComfyUIService {
    private baseUrl: string;

    constructor(baseUrl: string) {
        this.baseUrl = baseUrl;
    }

    async generateImage(prompt: string, settings: any) {
        try {
            // 构建 ComfyUI 工作流
            const workflow = this.buildWorkflow(prompt, settings);

            // 发送请求到 ComfyUI
            const response = await axios.post(`${this.baseUrl}/prompt`, workflow);

            // 等待生成完成
            const result = await this.waitForResult(response.data.prompt_id);

            return result;
        } catch (error) {
            console.error('ComfyUI generation error:', error);
            throw error;
        }
    }

    private buildWorkflow(prompt: string, settings: any) {
        // 这里需要根据 ComfyUI 的 API 构建正确的工作流
        return {
            "3": {
                "inputs": {
                    "text": prompt,
                    "clip": ["4", 0]
                },
                "class_type": "CLIPTextEncode"
            },
            "4": {
                "inputs": {
                    "ckpt_name": "v1-5-pruned.ckpt"
                },
                "class_type": "CheckpointLoaderSimple"
            },
            // ... 其他工作流节点
        };
    }

    private async waitForResult(promptId: string) {
        let retries = 0;
        const maxRetries = 30;

        while (retries < maxRetries) {
            try {
                const response = await axios.get(`${this.baseUrl}/history/${promptId}`);
                if (response.data.completed) {
                    // 获取生成的图片 URL
                    const imageUrl = `${this.baseUrl}/view?filename=${response.data.outputs[0].images[0].filename}`;
                    return { imageUrl };
                }
            } catch (error) {
                console.error('Error checking result:', error);
            }

            await new Promise(resolve => setTimeout(resolve, 1000));
            retries++;
        }

        throw new Error('Generation timeout');
    }
}

export function setupComfyUI() {
    const comfyuiUrl = process.env.COMFYUI_URL || 'http://localhost:8188';
    return new ComfyUIService(comfyuiUrl);
} 